import numpy as np
import os
import netCDF4 as nc
import proplot as plt
import cmaps
import gc
import scipy.ndimage as ndi
import random as rd



date='20250701'
labeldir=r'E:\akesu_radar\cheng\\'+date+'_label'
raddir=os.path.join(r'E:\akesu_radar\cheng',date)
outdir=os.path.join(r'E:\akesu_radar\cheng\pic',date)
if not os.path.exists(outdir):
    os.makedirs(outdir)

outtimes=[x.split('_')[0] for x in os.listdir(outdir) if 'png' in x]
labels=sorted([os.path.join(labeldir,x) for x in os.listdir(labeldir)])
rads=sorted([os.path.join(raddir,x) for x in os.listdir(raddir)])
nt=len(labels)

for it in range(nt):
    time=labels[it].split('\\')[-1].split('_')[0]
    print(time)
    if time in outtimes:
        continue
    ref=np.load(rads[it])['CR']
    refmax=np.where((ref>100)|(ref<0),np.nan,ref)

    ref_1_0=np.where(refmax>35,1,0)
    l=np.load(labels[it])


    sets=set(l.flatten())
    setlist=[x for x in sets]  # it is str
    setlist.remove(0)



    locs=ndi.center_of_mass(ref_1_0,l,setlist)


    fig,axs=plt.subplots()
    h=axs[0].contourf(refmax,levels=np.arange(0,80,5),cmap=cmaps.radar)
    axs[0].colorbar(h)

    # h2=axs[1].contourf(l,levels=np.arange(0,20,1),cmap='jet')
    # axs[1].colorbar(h2)

    for i,loc in enumerate(locs):
        jj,ii=np.where(l==setlist[i])
        jmax=jj.max()
        jmin=jj.min()
        imax=ii.max()
        imin=ii.min()
        print(imin,imax,jmin,jmax)
        axs[0].plot([imin,imin],[jmin,jmax],c='black',lw=0.5)
        axs[0].plot([imin,imax],[jmax,jmax],c='black',lw=0.5)
        axs[0].plot([imax, imax], [jmin, jmax], c='black',lw=0.5)
        axs[0].plot([imin, imax], [jmin, jmin], c='black',lw=0.5)
        axs[0].text(loc[1],loc[0],s=str(setlist[i]),size=6,color='black')

    picname=os.path.join(outdir,time+'_dbz_label.png')
    fig.canvas.print_figure(picname,dpi=400)
    plt.close()
    gc.collect()




